#include <aarch64.h>
#include <stdio.h>
#include "ft_assert.h"
#include <kernel.h>
#include "mmu.h"

#define assert FT_ASSERTNONERETURN
/******************************
 *  mmu config define
 ******************************/
#define CONFIG_ARM64_VA_BITS 42
#define CONFIG_ARM64_PA_BITS 42
#define CONFIG_MAX_XLAT_TABLES 8
/******************************
 *  mmu config define
 ******************************/

// extern const struct arm_mmu_region *mmu_regions;
// extern const uint32_t mmu_regions_size;
extern const struct arm_mmu_config mmu_config;

/* We support only 4kB translation granule */
#define PAGE_SIZE_SHIFT 12U
#define PAGE_SIZE (1U << PAGE_SIZE_SHIFT)
#define XLAT_TABLE_SIZE_SHIFT PAGE_SIZE_SHIFT /* Size of one complete table */
#define XLAT_TABLE_SIZE (1U << XLAT_TABLE_SIZE_SHIFT)

#define XLAT_TABLE_ENTRY_SIZE_SHIFT 3U /* Each table entry is 8 bytes */
#define XLAT_TABLE_LEVEL_MAX 3U

#define XLAT_TABLE_ENTRIES_SHIFT \
	(XLAT_TABLE_SIZE_SHIFT - XLAT_TABLE_ENTRY_SIZE_SHIFT)
#define XLAT_TABLE_ENTRIES (1U << XLAT_TABLE_ENTRIES_SHIFT)

/* Address size covered by each entry at given translation table level */
#define L3_XLAT_VA_SIZE_SHIFT PAGE_SIZE_SHIFT
#define L2_XLAT_VA_SIZE_SHIFT \
	(L3_XLAT_VA_SIZE_SHIFT + XLAT_TABLE_ENTRIES_SHIFT)
#define L1_XLAT_VA_SIZE_SHIFT \
	(L2_XLAT_VA_SIZE_SHIFT + XLAT_TABLE_ENTRIES_SHIFT)
#define L0_XLAT_VA_SIZE_SHIFT \
	(L1_XLAT_VA_SIZE_SHIFT + XLAT_TABLE_ENTRIES_SHIFT)

#define LEVEL_TO_VA_SIZE_SHIFT(level)              \
	(PAGE_SIZE_SHIFT + (XLAT_TABLE_ENTRIES_SHIFT * \
						(XLAT_TABLE_LEVEL_MAX - (level))))

/* Virtual Address Index within given translation table level */
#define XLAT_TABLE_VA_IDX(va_addr, level) \
	((va_addr >> LEVEL_TO_VA_SIZE_SHIFT(level)) & (XLAT_TABLE_ENTRIES - 1))

/*
 * Calculate the initial translation table level from CONFIG_ARM64_VA_BITS
 * For a 4 KB page size,
 * (va_bits <= 21)	 - base level 3
 * (22 <= va_bits <= 30) - base level 2
 * (31 <= va_bits <= 39) - base level 1
 * (40 <= va_bits <= 48) - base level 0
 */
#define GET_XLAT_TABLE_BASE_LEVEL(va_bits) \
	((va_bits > L0_XLAT_VA_SIZE_SHIFT)     \
		 ? 0U                              \
	 : (va_bits > L1_XLAT_VA_SIZE_SHIFT)   \
		 ? 1U                              \
	 : (va_bits > L2_XLAT_VA_SIZE_SHIFT)   \
		 ? 2U                              \
		 : 3U)

#define XLAT_TABLE_BASE_LEVEL GET_XLAT_TABLE_BASE_LEVEL(CONFIG_ARM64_VA_BITS)

#define GET_NUM_BASE_LEVEL_ENTRIES(va_bits) \
	(1U << (va_bits - LEVEL_TO_VA_SIZE_SHIFT(XLAT_TABLE_BASE_LEVEL)))

#define NUM_BASE_LEVEL_ENTRIES GET_NUM_BASE_LEVEL_ENTRIES(CONFIG_ARM64_VA_BITS)

static u64 base_xlat_table[NUM_BASE_LEVEL_ENTRIES] __aligned(NUM_BASE_LEVEL_ENTRIES * sizeof(u64));

static u64 xlat_tables[CONFIG_MAX_XLAT_TABLES][XLAT_TABLE_ENTRIES] __aligned(XLAT_TABLE_ENTRIES * sizeof(u64));

/* Translation table control register settings */
static u64 get_tcr(int el)
{
	u64 tcr;
	u64 pa_bits = CONFIG_ARM64_PA_BITS;
	u64 va_bits = CONFIG_ARM64_VA_BITS;
	u64 tcr_ps_bits;

	switch (pa_bits)
	{
	case 48:
		tcr_ps_bits = TCR_PS_BITS_256TB;
		break;
	case 44:
		tcr_ps_bits = TCR_PS_BITS_16TB;
		break;
	case 42:
		tcr_ps_bits = TCR_PS_BITS_4TB;
		break;
	case 40:
		tcr_ps_bits = TCR_PS_BITS_1TB;
		break;
	case 36:
		tcr_ps_bits = TCR_PS_BITS_64GB;
		break;
	default:
		tcr_ps_bits = TCR_PS_BITS_4GB;
		break;
	}

	if (el == 1)
	{
		tcr = (tcr_ps_bits << TCR_EL1_IPS_SHIFT);
		/*
		 * TCR_EL1.EPD1: Disable translation table walk for addresses
		 * that are translated using TTBR1_EL1.
		 */
		tcr |= TCR_EPD1_DISABLE;
	}
	else
		tcr = (tcr_ps_bits << TCR_EL3_PS_SHIFT);

	tcr |= TCR_T0SZ(va_bits);
	/*
	 * Translation table walk is cacheable, inner/outer WBWA and
	 * inner shareable
	 */
	tcr |= TCR_TG0_4K | TCR_SHARED_INNER | TCR_ORGN_WBWA | TCR_IRGN_WBWA;

	return tcr;
}

static int PteDescType(u64 *pte)
{
	return *pte & PTE_DESC_TYPE_MASK;
}

static u64 *CalculatePteIndex(u64 addr, unsigned int level)
{
	u32 base_level = XLAT_TABLE_BASE_LEVEL;
	u64 *pte;
	u64 idx;
	unsigned int i;

	/* Walk through all translation tables to find pte index */
	pte = (u64 *)base_xlat_table;
	for (i = base_level; i <= XLAT_TABLE_LEVEL_MAX; i++)
	{
		idx = XLAT_TABLE_VA_IDX(addr, i);
		pte += idx;
		/* Found pte index */
		if (i == level)
			return pte;
		/* if PTE is not table desc, can't traverse */
		if (PteDescType(pte) != PTE_TABLE_DESC)
			return NULL;
		/* Move to the next translation table level */
		pte = (u64 *)(*pte & 0x0000fffffffff000ULL);
	}

	return NULL;
}

static void SetPteTableDesc(u64 *pte, u64 *table, unsigned int level)
{
	(void)level;
	/* Point pte to new table */
	*pte = PTE_TABLE_DESC | (u64)table;
}

static void SetPteBlockDesc(u64 *pte, u64 addr_pa,
							unsigned int attrs, unsigned int level)
{
	u64 desc = addr_pa;
	unsigned int mem_type;

	desc |= (level == 3) ? PTE_PAGE_DESC : PTE_BLOCK_DESC;

	/* NS bit for security memory access from secure state */
	desc |= (attrs & MT_NS) ? PTE_BLOCK_DESC_NS : 0;

	/* AP bits for Data access permission */
	desc |= (attrs & MT_RW) ? PTE_BLOCK_DESC_AP_RW : PTE_BLOCK_DESC_AP_RO;

	/* the access flag */
	desc |= PTE_BLOCK_DESC_AF;

	/* memory attribute index field */
	mem_type = MT_TYPE(attrs);
	desc |= PTE_BLOCK_DESC_MEMTYPE(mem_type);

	switch (mem_type)
	{
	case MT_DEVICE_nGnRnE:
	case MT_DEVICE_nGnRE:
	case MT_DEVICE_GRE:
		/* Access to Device memory and non-cacheable memory are coherent
		 * for all observers in the system and are treated as
		 * Outer shareable, so, for these 2 types of memory,
		 * it is not strictly needed to set shareability field
		 */
		desc |= PTE_BLOCK_DESC_OUTER_SHARE;
		/* Map device memory as execute-never */
		desc |= PTE_BLOCK_DESC_PXN;
		desc |= PTE_BLOCK_DESC_UXN;
		break;
	case MT_NORMAL_NC:
	case MT_NORMAL:
		if (mem_type == MT_NORMAL)
			desc |= PTE_BLOCK_DESC_INNER_SHARE;
		else
			desc |= PTE_BLOCK_DESC_OUTER_SHARE;
	}

	*pte = desc;
}

/* Returns a new reallocated table */
static u64 *NewPreallocTable(void)
{
	static unsigned int table_idx;

	assert(table_idx < CONFIG_MAX_XLAT_TABLES);

	return (u64 *)(xlat_tables[table_idx++]);
}

/* Splits a block into table with entries spanning the old block */
static void SplitPteBlockDesc(u64 *pte, int level)
{
	u64 old_block_desc = *pte;
	u64 *new_table;
	unsigned int i = 0;
	/* get address size shift bits for next level */
	int levelshift = LEVEL_TO_VA_SIZE_SHIFT(level + 1);

	new_table = NewPreallocTable();

	for (i = 0; i < XLAT_TABLE_ENTRIES; i++)
	{
		new_table[i] = old_block_desc | (i << levelshift);

		if ((level + 1) == 3)
			new_table[i] |= PTE_PAGE_DESC;
	}

	/* Overwrite existing PTE set the new table into effect */
	SetPteTableDesc(pte, new_table, level);
}

/* Create/Populate translation table(s) for given region */
static void InitXlatTables(const struct arm_mmu_region *region)
{
	u64 *pte;
	u64 virt = region->base_va;
	u64 phys = region->base_pa;
	u64 size = region->size;
	u64 attrs = region->attrs;
	u64 level_size;
	u64 *new_table;
	unsigned int level = XLAT_TABLE_BASE_LEVEL;

	/* check minimum alignment requirement for given mmap region */
	assert(((virt & (PAGE_SIZE - 1)) == 0) &&
		   ((size & (PAGE_SIZE - 1)) == 0));

	while (size)
	{
		assert(level <= XLAT_TABLE_LEVEL_MAX);

		/* Locate PTE for given virtual address and page table level */
		pte = CalculatePteIndex(virt, level);
		assert(pte != NULL);

		level_size = 1ULL << LEVEL_TO_VA_SIZE_SHIFT(level);

		if (size >= level_size && !(virt & (level_size - 1)))
		{
			/* Given range fits into level size,
			 * create block/page descriptor
			 */
			SetPteBlockDesc(pte, phys, attrs, level);
			virt += level_size;
			phys += level_size;
			size -= level_size;
			/* Range is mapped, start again for next range */
			level = XLAT_TABLE_BASE_LEVEL;
		}
		else if (PteDescType(pte) == PTE_INVALID_DESC)
		{
			/* Range doesn't fit, create subtable */
			new_table = NewPreallocTable();
			SetPteTableDesc(pte, new_table, level);
			level++;
		}
		else if (PteDescType(pte) == PTE_BLOCK_DESC)
		{
			SplitPteBlockDesc(pte, level);
			level++;
		}
		else if (PteDescType(pte) == PTE_TABLE_DESC)
			level++;
	}
}

static void SetupPageTables(void)
{
	unsigned int index;
	const struct arm_mmu_region *region;
	u64 max_va = 0, max_pa = 0;

	for (index = 0; index < mmu_config.num_regions; index++)
	{
		region = &mmu_config.mmu_regions[index];
		max_va = max(max_va, region->base_va + region->size);
		max_pa = max(max_pa, region->base_pa + region->size);
	}

	assert(max_va <= (1ULL << CONFIG_ARM64_VA_BITS));
	assert(max_pa <= (1ULL << CONFIG_ARM64_PA_BITS));

	/* create translation tables for user provided platform regions */
	for (index = 0; index < mmu_config.num_regions; index++)
	{
		region = &mmu_config.mmu_regions[index];
		if (region->size || region->attrs)
			InitXlatTables(region);
	}
}

static void EnableMmuEl1()
{

	/* Set MAIR, TCR and TBBR registers */
	__asm__ volatile("msr mair_el1, %0"
					 :
					 : "r"(MEMORY_ATTRIBUTES)
					 : "memory", "cc");
	__asm__ volatile("msr tcr_el1, %0"
					 :
					 : "r"(get_tcr(1))
					 : "memory", "cc");
	__asm__ volatile("msr ttbr0_el1, %0"
					 :
					 : "r"((u64)base_xlat_table)
					 : "memory", "cc");

	/* Ensure these changes are seen before MMU is enabled */
	ISB();
}

/* ARM MMU Driver Initial Setup */

/*
 * @brief MMU default configuration
 *
 * This function provides the default configuration mechanism for the Memory
 * Management Unit (MMU).
 */
int MmuInit()
{
	u64 val;

	/* Current MMU code supports only EL1 */
	__asm__ volatile("mrs %0, CurrentEL"
					 : "=r"(val));

	assert(GET_EL(val) == MODE_EL1);

	/* Ensure that MMU is already not enabled */
	__asm__ volatile("mrs %0, sctlr_el1"
					 : "=r"(val));
	assert((val & SCTLR_ELx_M) == 0);

	SetupPageTables();

	/* currently only EL1 is supported */
	EnableMmuEl1();

	return 0;
}
